<!DOCTYPE html>
<html>

<head>
    <title>WebGPU</title>
    <style>
        body {
            margin: 0;
            overflow: hidden;
        }

        canvas {
            display: block;
        }
    </style>
</head>

<body>

    <canvas id="webgpu-canvas"></canvas>
    <script type="module">
        // const PARTICLE_COUNT = 400;
        import PARTICLE_COUNT from "./config.js"

        async function initWebGPU() {
            const canvas = document.getElementById('webgpu-canvas');
            canvas.width = window.innerWidth;
            canvas.height = window.innerHeight;

            const adapter = await navigator.gpu.requestAdapter();
            const device = await adapter.requestDevice();
            const context = canvas.getContext('webgpu');
            const format = navigator.gpu.getPreferredCanvasFormat();
            context.configure({ device, format });

            // 创建粒子初始数据
            const initialParticles = new Float32Array(PARTICLE_COUNT * 4); // x, y, vx, vy
            for (let i = 0; i < PARTICLE_COUNT * 4; i += 4) {
                initialParticles[i] = Math.random() * 2 - 1; // x
                initialParticles[i + 1] = Math.random() * 2 - 1; // y
                initialParticles[i + 2] = (Math.random() - 0.5) * 0.01; // vx
                initialParticles[i + 3] = (Math.random() - 0.5) * 0.01; // vy
            }

            // 创建存储缓冲区
            const particleBuffer = device.createBuffer({
                size: initialParticles.byteLength,
                usage: GPUBufferUsage.STORAGE | GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
                mappedAtCreation: true
            });
            new Float32Array(particleBuffer.getMappedRange()).set(initialParticles);
            particleBuffer.unmap();

            // 计算管线
            const computeShader = device.createShaderModule({
                code: `
            @group(0) @binding(0) var<storage, read_write> particles : array<vec4f>;
    
            @compute @workgroup_size(64)
            fn main(@builtin(global_invocation_id) id: vec3u) {
                var p = particles[id.x];
                p.x += p.z;
                p.y += p.w;
    
                if (abs(p.x) >= 1.0) { p.z *= -1.0; }
                if (abs(p.y) >= 1.0) { p.w *= -1.0; }
    
                particles[id.x] = p;
            }`
            });

            const computePipeline = device.createComputePipeline({
                layout: 'auto',
                compute: { module: computeShader, entryPoint: 'main' }
            });

            // 渲染管线
            const renderPipeline = device.createRenderPipeline({
                layout: 'auto',
                vertex: {
                    module: device.createShaderModule({
                        code: `
                    @vertex
                    fn main(@location(0) pos: vec2f) -> @builtin(position) vec4f {
                        return vec4f(pos, 1.0, 1.0);
                    }`
                    }),
                    entryPoint: 'main',
                    buffers: [{
                        arrayStride: 4 * 4, // x, y, vx, vy
                        attributes: [{ shaderLocation: 0, offset: 0, format: 'float32x2' }]
                    }]
                },
                fragment: {
                    module: device.createShaderModule({
                        code: `
                    @fragment
                    fn main() -> @location(0) vec4f {
                        return vec4f(1.0);
                    }`
                    }),
                    entryPoint: 'main',
                    targets: [{ format }]
                },
                primitive: { topology: 'point-list' }
            });

            // 绑定组
            const bindGroup = device.createBindGroup({
                layout: computePipeline.getBindGroupLayout(0),
                entries: [{ binding: 0, resource: { buffer: particleBuffer } }]
            });

            // 渲染循环
            function animate() {
                const encoder = device.createCommandEncoder();

                // 计算通道
                const computePass = encoder.beginComputePass();
                computePass.setPipeline(computePipeline);
                computePass.setBindGroup(0, bindGroup);
                computePass.dispatchWorkgroups(Math.ceil(PARTICLE_COUNT / 64));
                computePass.end();

                // 渲染通道
                const renderPass = encoder.beginRenderPass({
                    colorAttachments: [{
                        view: context.getCurrentTexture().createView(),
                        loadOp: 'clear',
                        storeOp: 'store'
                    }]
                });
                renderPass.setPipeline(renderPipeline);
                renderPass.setVertexBuffer(0, particleBuffer);
                renderPass.draw(PARTICLE_COUNT);
                renderPass.end();

                device.queue.submit([encoder.finish()]);
                requestAnimationFrame(animate);
            }
            animate();
        }
        initWebGPU();
    </script>
</body>

</html>